# this file is used to describe the information about station
import pandas as pd


class FireStation:
    def __init__(self, node_name: str, station_name: str, max_output: float, upper_f1, upper_f2, low_f1, low_f2,
                 restart_cost: float, cost_f1, cost_f2, cost_f3):
        self.NodeName = node_name
        self.StationName: station_name
        self.OutputRange = [max_output * low_f1, max_output *
                            low_f2, max_output * upper_f2, max_output * upper_f1]
        self.RestartCost = restart_cost
        self.a = cost_f1
        self.b = cost_f2
        self.c = cost_f3

    def step(self, action: float, last_action: float) -> float:
        r = 0.0
        if (action == 0 and last_action != 0) or (last_action == 0 and action != 0):
            r -= self.RestartCost
        # check action is in range
        # if action < self.OutputRange[0]:
        #     r -= self.OutputRange[0] - action
        # if action > self.OutputRange[3]:
        #     r -= action - self.OutputRange[3]
        # calculate the cost of coals
        coals = self.a * action ** 2 + self.b * action + self.c
        r -= coals
        # return r * 1e-2
        return 0


class NewEnergyStation:
    def __init__(self, predict_output, actual_output, installed_capacity):
        self.PredictOutput = predict_output
        self.ActualOutput = actual_output
        self.InstalledCapacity = installed_capacity

    def step(self, action, curr_t):
        r = 0
        if action > self.ActualOutput[curr_t]:
            r -= action - self.ActualOutput[curr_t]
        else:
            r += action
        return r


class WaterStation:
    def __init__(self, node_name, station_name, output):
        # todo: water station we need to consider
        pass


class StorageStation:
    def __init__(self, node_name, station_name, storage_type: str, output, hours, hours_in_year, rate):
        self.NodeName = node_name
        self.StationName = station_name
        self.StorageType = storage_type
        self.Output = output
        self.MaxCapacity = self.Output * hours
        self.HourLimitInYear = hours_in_year
        self.Rate = rate
        self.Capacity = 0

    def step(self, action):
        r = 0
        if action < 0:
            # charge storage
            self.Capacity -= action * self.Rate
            if self.Capacity > self.MaxCapacity:
                r -= self.Capacity - self.MaxCapacity
                self.Capacity = self.MaxCapacity
        else:
            # discharge storage
            # action = abs(action)
            self.Capacity -= action * self.Rate
            if self.Capacity < 0:
                r -= 0 - self.Capacity
                self.Capacity = 0
        return r*1000
